"""
# Swift-Hohenberg Equation - Advanced Pattern Formation

This module implements the Swift-Hohenberg equation, a fundamental model for 
pattern formation near convective instability. It's widely used in studying
convection patterns, Taylor-Couette flow, and general pattern formation theory.

## Mathematical Background

The Swift-Hohenberg equation is:
```
∂u/∂t = r·u - (∇² + k₀²)²u - u³
```

Where:
- u: Order parameter (e.g., convection amplitude)
- r: Control parameter (distance from instability threshold)
- k₀: Critical wavenumber
- (∇² + k₀²)²: Fourth-order differential operator

## Key Features

- Linear instability at r = 0 leads to patterns with wavenumber k₀
- Subcritical (r < 0) and supercritical (r > 0) regimes  
- Rich dynamics including stripes, hexagons, and localized structures
- Amplitude equations for weakly nonlinear analysis

## Applications

- Rayleigh-Bénard convection
- Taylor-Couette flow between cylinders
- Faraday waves on liquid surfaces
- General pattern formation theory

## References

- Swift, J. & Hohenberg, P. C. (1977). Phys. Rev. A 15, 319.
- Cross, M. C. & Hohenberg, P. C. (1993). Rev. Mod. Phys. 65, 851.
- Hoyle, R. (2006). Pattern Formation: An Introduction to Methods.
"""

using LinearAlgebra, FFTW, Plots, Statistics, Printf

export SwiftHohenbergSystem, run_swift_hohenberg, analyze_sh_patterns

"""
    SwiftHohenbergSystem

Configuration for Swift-Hohenberg equation simulations.

# Fields
- `r::Float64`: Control parameter (bifurcation parameter)
- `k0::Float64`: Critical wavenumber
- `domain_size::Tuple`: Spatial domain (Lx, Ly)
- `grid_size::Tuple`: Grid resolution (nx, ny)
- `boundary_condition::Symbol`: Boundary conditions (:periodic, :dirichlet)
"""
struct SwiftHohenbergSystem
    r::Float64
    k0::Float64
    domain_size::Tuple{Float64, Float64}
    grid_size::Tuple{Int, Int}
    boundary_condition::Symbol
end

"""
    SwiftHohenbergSystem(r, k0=1.0; domain_size=(20.0, 20.0), grid_size=(128, 128), 
                        boundary_condition=:periodic)

Create Swift-Hohenberg system configuration.

# Arguments
- `r`: Control parameter (r > 0 for supercritical, r < 0 for subcritical)
- `k0`: Critical wavenumber (determines pattern wavelength λ = 2π/k0)
"""
function SwiftHohenbergSystem(r, k0=1.0; 
                             domain_size=(20.0, 20.0), 
                             grid_size=(128, 128),
                             boundary_condition=:periodic)
    return SwiftHohenbergSystem(r, k0, domain_size, grid_size, boundary_condition)
end

"""
    create_differential_operator(system::SwiftHohenbergSystem)

Create the fourth-order differential operator (∇² + k₀²)² using FFT.
"""
function create_differential_operator(system::SwiftHohenbergSystem)
    nx, ny = system.grid_size
    Lx, Ly = system.domain_size
    
    # Wavenumbers for FFT
    kx = 2π * fftfreq(nx, nx/Lx)
    ky = 2π * fftfreq(ny, ny/Ly)
    
    # Create wavenumber grids
    Kx = kx' .* ones(ny)
    Ky = ones(nx)' .* ky
    K2 = Kx.^2 + Ky.^2
    
    # Fourth-order operator: (∇² + k₀²)² = (k² + k₀²)²
    operator = -(K2 .+ system.k0^2).^2
    
    return operator, K2
end

"""
    initialize_swift_hohenberg(system::SwiftHohenbergSystem, 
                              initial_type::Symbol=:random_noise)

Initialize the order parameter field with appropriate initial conditions.

# Arguments
- `initial_type::Symbol`: Type of initial condition
  - `:random_noise`: Small random perturbations
  - `:stripe_pattern`: Initial stripe pattern
  - `:hexagon_pattern`: Hexagonal pattern seed
  - `:localized_pulse`: Localized structure
"""
function initialize_swift_hohenberg(system::SwiftHohenbergSystem, 
                                   initial_type::Symbol=:random_noise)
    nx, ny = system.grid_size
    Lx, Ly = system.domain_size
    
    # Create coordinate grids
    x = range(0, Lx, length=nx)
    y = range(0, Ly, length=ny)
    X = [xi for xi in x, yi in y]  
    Y = [yi for xi in x, yi in y]
    
    if initial_type == :random_noise
        # Small random perturbations
        u = 0.1 * randn(nx, ny)
        
    elseif initial_type == :stripe_pattern
        # Stripe pattern with critical wavenumber
        u = 0.1 * cos.(system.k0 * X)
        u .+= 0.01 * randn(nx, ny)  # Add small noise
        
    elseif initial_type == :hexagon_pattern
        # Hexagonal pattern (superposition of three stripe modes)
        k = system.k0
        u = 0.1 * (cos.(k * X) + 
                  cos.(k * (0.5*X - √3/2*Y)) + 
                  cos.(k * (0.5*X + √3/2*Y)))
        u .+= 0.01 * randn(nx, ny)
        
    elseif initial_type == :localized_pulse
        # Gaussian pulse for studying localized structures
        x0, y0 = Lx/2, Ly/2
        σ = 2π/system.k0  # Width related to pattern wavelength
        u = 0.5 * exp.(-((X .- x0).^2 + (Y .- y0).^2) / (2σ^2))
        
    else
        throw(ArgumentError("Unknown initial type: $initial_type"))
    end
    
    return u, X, Y
end

"""
    evolve_swift_hohenberg!(u, u_fft, operator, system::SwiftHohenbergSystem, dt)

Perform one time step using pseudo-spectral methods (FFT).
"""
function evolve_swift_hohenberg!(u, u_fft, operator, system::SwiftHohenbergSystem, dt)
    # Forward FFT
    u_fft .= fft(u)
    
    # Linear part in Fourier space: ∂u/∂t = r·u - (∇² + k₀²)²u
    linear_part = system.r .+ operator
    
    # Nonlinear part: -u³ (computed in real space)
    nonlinear = -u.^3
    nonlinear_fft = fft(nonlinear)
    
    # Time stepping using exponential time differencing (ETD)
    # For simplicity, use first-order ETD: u(t+dt) = exp(L·dt)u(t) + ∫₀ᵈᵗ exp(L(dt-s))N(s)ds
    exp_linear = exp.(linear_part * dt)
    
    # First-order approximation for nonlinear term
    u_fft .= exp_linear .* u_fft + dt * exp_linear .* nonlinear_fft
    
    # Inverse FFT
    u .= real.(ifft(u_fft))
end

"""
    run_swift_hohenberg(system::SwiftHohenbergSystem;
                       t_final=100.0, dt=nothing, 
                       initial_type=:random_noise,
                       save_interval=1.0, verbose=true)

Run Swift-Hohenberg equation simulation.
"""
function run_swift_hohenberg(system::SwiftHohenbergSystem;
                            t_final=100.0, dt=nothing,
                            initial_type=:random_noise,
                            save_interval=1.0, verbose=true)
    
    if verbose
        println("🌊 Swift-Hohenberg Pattern Formation")
        println("===================================")
        println("Control parameter r = $(system.r)")
        println("Critical wavenumber k₀ = $(system.k0)")
        println("Pattern wavelength λ = $(round(2π/system.k0, digits=2))")
        println("Domain: $(system.domain_size)")
        println("Grid: $(system.grid_size)")
        println()
        
        # Stability information
        if system.r > 0
            println("Regime: Supercritical (pattern-forming)")
        else
            println("Regime: Subcritical (pattern-suppressing)")
        end
        println()
    end
    
    # Initialize
    u, X, Y = initialize_swift_hohenberg(system, initial_type)
    nx, ny = system.grid_size
    
    # Create differential operator
    operator, K2 = create_differential_operator(system)
    
    # Determine stable time step
    if dt === nothing
        # CFL-like condition for fourth-order operator
        max_eigenvalue = maximum(abs.(system.r .+ operator))
        dt = 0.01 / max_eigenvalue
        dt = min(dt, 0.01)  # Upper bound for stability
    end
    
    n_steps = Int(t_final / dt)
    n_save = max(1, Int(save_interval / dt))
    
    if verbose
        println("Time integration:")
        println("  dt = $(round(dt, digits=6))")
        println("  Steps: $n_steps")
        println("  Save interval: $(save_interval)")
        println()
    end
    
    # Storage
    times = Float64[]
    snapshots = Array{Float64,3}(undef, nx, ny, 0)
    energy_history = Float64[]
    
    # Work arrays
    u_fft = similar(u, Complex{Float64})
    
    # Save initial condition
    t = 0.0
    push!(times, t)
    snapshots = cat(snapshots, reshape(u, nx, ny, 1), dims=3)
    
    # Calculate initial energy
    energy = calculate_energy(u, K2, system)
    push!(energy_history, energy)
    
    if verbose
        print("Progress: ")
    end
    
    # Time evolution
    for step in 1:n_steps
        # Time step
        evolve_swift_hohenberg!(u, u_fft, operator, system, dt)
        t += dt
        
        # Save data
        if step % n_save == 0
            push!(times, t)
            snapshots = cat(snapshots, reshape(u, nx, ny, 1), dims=3)
            
            # Calculate energy
            energy = calculate_energy(u, K2, system)
            push!(energy_history, energy)
            
            if verbose
                progress = step / n_steps * 100
                print("$(Int(round(progress)))% ")
            end
        end
    end
    
    if verbose
        println("\n✅ Simulation completed!")
        println()
    end
    
    return (
        times = times,
        u_final = u,
        snapshots = snapshots,
        energy_history = energy_history,
        X = X,
        Y = Y,
        system = system,
        dt = dt
    )
end

"""
    calculate_energy(u, K2, system::SwiftHohenbergSystem)

Calculate the Swift-Hohenberg free energy functional.
"""
function calculate_energy(u, K2, system::SwiftHohenbergSystem)
    nx, ny = system.grid_size
    dx = system.domain_size[1] / nx
    dy = system.domain_size[2] / ny
    
    # Energy density in Fourier space for quadratic terms
    u_fft = fft(u)
    
    # Quadratic part: ∫ u(-r + (∇² + k₀²)²)u dx
    quadratic_energy = real(sum(conj.(u_fft) .* (-system.r .+ (K2 .+ system.k0^2).^2) .* u_fft))
    quadratic_energy *= dx * dy / (nx * ny)
    
    # Quartic part: ∫ u⁴/4 dx  
    quartic_energy = sum(u.^4) * dx * dy / 4
    
    total_energy = quadratic_energy + quartic_energy
    return total_energy
end

"""
    analyze_sh_patterns(results; verbose=true)

Analyze patterns in Swift-Hohenberg simulation results.
"""
function analyze_sh_patterns(results; verbose=true)
    if verbose
        println("📊 Swift-Hohenberg Pattern Analysis")
        println("===================================")
    end
    
    u_final = results.u_final
    system = results.system
    
    # Basic statistics
    u_mean = mean(u_final)
    u_std = std(u_final)
    u_max = maximum(u_final)
    u_min = minimum(u_final)
    
    if verbose
        println("Field statistics:")
        println("  Mean: $(round(u_mean, digits=4))")
        println("  Std:  $(round(u_std, digits=4))")
        println("  Range: [$(round(u_min, digits=4)), $(round(u_max, digits=4))]")
    end
    
    # Fourier analysis
    u_centered = u_final .- u_mean
    u_fft = fft(u_centered)
    power_spectrum = abs.(u_fft).^2
    
    # Wavenumber grids
    nx, ny = size(u_final)
    kx = 2π * fftfreq(nx, nx/system.domain_size[1])
    ky = 2π * fftfreq(ny, ny/system.domain_size[2])
    K_magnitude = sqrt.(kx'.^2 .+ ky.^2)
    
    # Find peak in power spectrum
    # Exclude DC component
    power_no_dc = copy(power_spectrum)
    power_no_dc[1, 1] = 0
    
    peak_idx = argmax(power_no_dc)
    peak_k = K_magnitude[peak_idx]
    peak_wavelength = peak_k > 0 ? 2π / peak_k : Inf
    
    # Compare with theoretical prediction
    theoretical_k = system.k0
    theoretical_wavelength = 2π / theoretical_k
    
    if verbose
        println("Spectral analysis:")
        println("  Peak wavenumber: $(round(peak_k, digits=4))")
        println("  Peak wavelength: $(round(peak_wavelength, digits=2))")
        println("  Theoretical k₀:  $(round(theoretical_k, digits=4))")
        println("  Theoretical λ₀:  $(round(theoretical_wavelength, digits=2))")
        
        if theoretical_k > 0
            k_error = abs(peak_k - theoretical_k) / theoretical_k * 100
            println("  Wavenumber error: $(round(k_error, digits=1))%")
        end
    end
    
    # Pattern classification based on structure tensor
    # Calculate gradients
    ∇u_x, ∇u_y = gradient_finite_diff(u_final)
    
    # Structure tensor components  
    J11 = mean(∇u_x.^2)
    J12 = mean(∇u_x .* ∇u_y)
    J22 = mean(∇u_y.^2)
    
    # Eigenvalues of structure tensor
    trace = J11 + J22
    det = J11 * J22 - J12^2
    λ1 = (trace + sqrt(trace^2 - 4*det)) / 2
    λ2 = (trace - sqrt(trace^2 - 4*det)) / 2
    
    # Anisotropy measure
    if λ1 + λ2 > 0
        anisotropy = abs(λ1 - λ2) / (λ1 + λ2)
    else
        anisotropy = 0
    end
    
    # Pattern classification
    pattern_type = if u_std < 0.01
        "no_pattern"
    elseif anisotropy > 0.5
        "stripes"
    elseif anisotropy < 0.2 && u_std > 0.1
        "hexagons"
    else
        "mixed_pattern"
    end
    
    if verbose
        println("Pattern characterization:")
        println("  Type: $pattern_type") 
        println("  Anisotropy: $(round(anisotropy, digits=3))")
        println()
    end
    
    # Energy analysis
    if haskey(results, :energy_history)
        final_energy = results.energy_history[end]
        energy_change = results.energy_history[end] - results.energy_history[1]
        
        if verbose
            println("Energy analysis:")
            println("  Final energy: $(round(final_energy, digits=4))")
            println("  Energy change: $(round(energy_change, digits=4))")
        end
    end
    
    return (
        statistics = (mean=u_mean, std=u_std, min=u_min, max=u_max),
        spectral = (peak_k=peak_k, peak_wavelength=peak_wavelength,
                   theoretical_k=theoretical_k, theoretical_wavelength=theoretical_wavelength),
        pattern_type = pattern_type,
        anisotropy = anisotropy,
        power_spectrum = power_spectrum
    )
end

"""
    gradient_finite_diff(u)

Compute gradients using finite differences with periodic boundary conditions.
"""
function gradient_finite_diff(u)
    nx, ny = size(u)
    
    ∇u_x = zeros(nx, ny)
    ∇u_y = zeros(nx, ny)
    
    # Central differences with periodic boundaries
    for j in 1:ny, i in 1:nx
        i_prev = i == 1 ? nx : i-1
        i_next = i == nx ? 1 : i+1
        j_prev = j == 1 ? ny : j-1
        j_next = j == ny ? 1 : j+1
        
        ∇u_x[i,j] = (u[i_next,j] - u[i_prev,j]) / 2
        ∇u_y[i,j] = (u[i,j_next] - u[i,j_prev]) / 2
    end
    
    return ∇u_x, ∇u_y
end

"""
    bifurcation_diagram(k0=1.0; r_range=(-0.1, 0.1), n_points=20, kwargs...)

Generate bifurcation diagram showing pattern amplitude vs control parameter.
"""
function bifurcation_diagram(k0=1.0; r_range=(-0.1, 0.1), n_points=20, kwargs...)
    println("📈 Generating Bifurcation Diagram")
    println("=================================")
    
    r_values = range(r_range[1], r_range[2], length=n_points)
    amplitudes = Float64[]
    
    for (i, r) in enumerate(r_values)
        print("Point $i/$n_points (r = $(round(r, digits=4))): ")
        
        system = SwiftHohenbergSystem(r, k0; kwargs...)
        results = run_swift_hohenberg(system, t_final=50.0, verbose=false)
        analysis = analyze_sh_patterns(results, verbose=false)
        
        amplitude = analysis.statistics.std
        push!(amplitudes, amplitude)
        
        println("amplitude = $(round(amplitude, digits=4))")
    end
    
    # Create bifurcation plot
    p = plot(r_values, amplitudes, marker=:circle, linewidth=2,
             xlabel="Control Parameter r", ylabel="Pattern Amplitude",
             title="Swift-Hohenberg Bifurcation Diagram",
             legend=false, grid=true)
    
    # Add theoretical bifurcation point
    vline!([0.0], linestyle=:dash, color=:red, 
           label="Theoretical bifurcation", linewidth=2)
    
    println("\n✅ Bifurcation diagram completed!")
    return p, (r_values=collect(r_values), amplitudes=amplitudes)
end

# Demonstration examples
if abspath(PROGRAM_FILE) == @__FILE__
    println("🌊 Swift-Hohenberg Equation - Research Examples")
    println("=" * 50)
    
    # Example 1: Supercritical pattern formation
    println("\n1. Supercritical Pattern Formation (r > 0)")
    system1 = SwiftHohenbergSystem(0.05, 1.0)  # r > 0
    results1 = run_swift_hohenberg(system1, t_final=50.0)
    analysis1 = analyze_sh_patterns(results1)
    
    # Example 2: Subcritical regime
    println("\n2. Subcritical Regime (r < 0)")
    system2 = SwiftHohenbergSystem(-0.02, 1.0)  # r < 0
    results2 = run_swift_hohenberg(system2, t_final=50.0, 
                                  initial_type=:hexagon_pattern)
    analysis2 = analyze_sh_patterns(results2)
    
    # Example 3: Localized structures
    println("\n3. Localized Structures")
    system3 = SwiftHohenbergSystem(-0.01, 1.0, domain_size=(40.0, 40.0))
    results3 = run_swift_hohenberg(system3, t_final=100.0,
                                  initial_type=:localized_pulse)
    analysis3 = analyze_sh_patterns(results3)
    
    # Example 4: Bifurcation analysis
    println("\n4. Bifurcation Analysis")
    bif_plot, bif_data = bifurcation_diagram(1.0, r_range=(-0.05, 0.1), n_points=15)
    
    println("\n🎉 All Swift-Hohenberg examples completed!")
    println("Ready for advanced pattern formation research.")
end